/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "uhm/common.hxx"
#include "uhm/const.hxx"
#include "uhm/util.hxx"

#include "uhm/object.hxx"

#include "uhm/operation/scheduler.hxx"
#include "uhm/operation/element.hxx"

#include "uhm/mesh/element.hxx"

#include "uhm/matrix/uhm/matrix.hxx"

#include "uhm/mesh/mesh.hxx"

#ifdef UHM_INTERF_PETSC_ENABLE
#include "petscksp.h"
#endif

#include "uhm/interf/sparse.hxx"
#include "uhm/interf/petsc.hxx"

#define UHM_INTERF_PETSC_COPY

namespace uhm {

#undef  __FUNCT__
#define __FUNCT__ "Mesh_::export_matrix"
  bool Mesh_::export_matrix(interf::Petsc petsc, int n_rhs, int is_sym) {

    // sparse enviroments : fortran index, compressed row;
    int fort=0, cs=0;

    // sparse object
    interf::Sparse sp;

    // easy setting
    if (petsc->is_complex()) 
      sp = new interf::ZSparse_(fort, cs, n_rhs);
    else 
      sp = new interf::DSparse_(fort, cs, n_rhs);
    
    this->export_matrix(sp, n_rhs);
    
    petsc->create_sparse_matrix(sp, is_sym);

    printf("- PETSC INTERFACED -\n");
    printf("n_dof       : %d\n", (int)petsc->get_n_dof());
    printf("n_nonzero   : %d\n", (int)petsc->get_n_nonzero());
    printf("n_rhs       : %d\n", (int)petsc->get_n_rhs());
    printf("is complex  : %d\n", (int)petsc->is_complex());

    return true;
  }

  namespace interf {

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::export_matrix"
    bool Petsc_::export_matrix(Mesh m) {
      // sparse enviroments : fortran index, compressed row;
      int fort=0, cs=0;

      // sparse object
      Sparse sp;

      // easy setting
      if (this->is_complex())
        sp = new ZSparse_(fort, cs, n_rhs);
      else
        sp = new DSparse_(fort, cs, n_rhs);

      m->import_matrix(sp, UHM_ASSEMBLED, this->xx, this->n_dof, this->n_rhs);

      return true;
    }


  
#undef  __FUNCT__
#define __FUNCT__ "Petsc_::disp"
    void Petsc_::disp(FILE *stream) {
      // nothing yet
    }
  
    // ----------------------------------------------------------------
    // ** petsc control


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::is_complex"
    bool Petsc_::is_complex() { return (this->datatype == UHM_COMPLEX); }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::is_spd"
    bool Petsc_::is_spd()     { return this->spd; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::set_spd"
    void Petsc_::set_spd(bool is_spd)         { this->spd = is_spd; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::set_tolerence"
    void Petsc_::set_tolerence(double tol) { this->tol = tol; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::set_show_n_rhs"
    void Petsc_::set_show_n_rhs(int show_n_rhs) { this->show_n_rhs = show_n_rhs; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::create_sparse_matrix"
    int Petsc_::create_sparse_matrix(Sparse sp, int is_sym)   {
    
      if (this->created) 
        UHM_ERROR(false, ">> Already created");

      //     // too slow in interfacing to PETSC
      //     sp->triplet( is_sym,
      //                  this->n_dof, this->n_rhs, this->n_nz,
      //                  this->ia, this->ja,
      //                  this->aa, this->bb );
    
      sp->compress( is_sym,
                    this->n_dof, this->n_rhs, this->n_nz,
                    this->ia, this->ja,
                    this->aa, this->bb );

      this->xx.clear();
      this->xx.reserve( this->bb.size() );

      for (int k=0;k<this->bb.size();++k)
        this->xx.push_back( 0.0 );

      if (this->show_n_rhs > (int)this->bb.size())
        this->show_n_rhs = (int)this->bb.size();

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;

      // ** Matrix Interface
      //
      // - Options
      // -mat_type seqaij   - AIJ type, uses MatCreateSeqAIJ()
      // -mat_type mpiaij   - AIJ type, uses MatCreateMPIAIJ()
      // -mat_type seqdense - dense type, uses MatCreateSeqDense()
      // -mat_type mpidense - dense type, uses MatCreateMPIDense()
      // -mat_type seqbaij  - block AIJ type, uses MatCreateSeqBAIJ()
      // -mat_type mpibaij  - block AIJ type, uses MatCreateMPIBAIJ()

      // ** does not work don't know why
      // ierr = MatSeqAIJSetPreallocationCSR(this->A, 
      //                                    &(this->ia[0]), &(this->ja[0]), 
      //                                    &(this->aa[0]));CHKERRQ(ierr);

      // - Options
      // ADD_VALUES, INSERT_VALUES

      //     // IJ option SLOW....
      //     for (int k=0;k<aa.size();++k)
      //       ierr = MatSetValues(this->A, 
      //                           1, &(this->ia[k]), 
      //                           1, &(this->ja[k]),
      //                           &(this->aa[k]), INSERT_VALUES);CHKERRQ(ierr);

      std::vector<int> row_n_nz;
      row_n_nz.reserve( this->n_dof );
      for (int i=0;i<this->n_dof;++i) {
        int cnt=0;
        for (int j=this->ia[i];j<this->ia[i+1];++j)
          ++cnt;
        row_n_nz.push_back(cnt);
      }

      // follow three lines are same as MatCreateSeqAIJ
      ierr = MatCreate(PETSC_COMM_WORLD, &(this->A));CHKERRQ(ierr);
      ierr = MatSetSizes(this->A, PETSC_DECIDE,PETSC_DECIDE, 
                         this->n_dof, this->n_dof);CHKERRQ(ierr);
      ierr = MatSetType(this->A, MATSEQAIJ);CHKERRQ(ierr);

      ierr = MatSetFromOptions(this->A);CHKERRQ(ierr);

      // nz is ignored so set zero
      ierr = MatSeqAIJSetPreallocation(this->A,0,&(row_n_nz[0]));CHKERRQ(ierr);

      int stride = ( this->is_complex() ? 2 : 1 );
      for (int i=0;i<this->n_dof;++i) {
        int begin = this->ia[i], end = this->ia[i+1];

        ierr = MatSetValues(this->A,
                            1, &i,
                            end-begin, &(this->ja[begin]),
                            &(this->aa[begin*stride]),
                            INSERT_VALUES);CHKERRQ(ierr);
      }

      // check the structure
      ierr = MatSetOption(this->A, 
                          MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE);CHKERRQ(ierr);



      // - Options
      // MAT_FLUSH_ASSEMBLY - When it change insert to add and vice versa
      // MAT_FINAL_ASSEMBLY
      ierr = MatAssemblyBegin(this->A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      ierr = MatAssemblyEnd  (this->A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

      // aa is not necessary anymore
      this->aa.clear();
      this->ia.clear();
      this->ja.clear();

      // ** Vector Interface
    
      ierr = VecCreate(PETSC_COMM_WORLD, &this->b);CHKERRQ(ierr);
      ierr = VecSetSizes(this->b, PETSC_DECIDE, this->n_dof);CHKERRQ(ierr);
      ierr = VecSetFromOptions(this->b);CHKERRQ(ierr);
      ierr = VecDuplicate(this->b, &this->x);CHKERRQ(ierr);
      ierr = VecDuplicate(this->x, &this->r);CHKERRQ(ierr);
    
#endif

      for (int k=0; k<this->show_n_rhs; ++k)
        printf(" x[ %d ] = % E, b [ %d ] = % E\n",
               k, this->xx.at(k),
               k, this->bb.at(k) );

      this->created = true;

      return true;
    }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::delete_sparse_matrix"
    int Petsc_::delete_sparse_matrix() {
      if (!this->created) return true;

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;

      ierr = MatDestroy(&this->A);CHKERRQ(ierr);
      ierr = VecDestroy(&this->x);CHKERRQ(ierr);
      ierr = VecDestroy(&this->b);CHKERRQ(ierr);

#endif

      this->ia.clear();
      this->ja.clear();
    
      this->aa.clear();
      this->xx.clear();
      this->bb.clear();

      return true;
    }
  

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::get_n_dof"
    int  Petsc_::get_n_dof()     { return this->n_dof; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::get_n_rhs"
    int  Petsc_::get_n_rhs()     { return this->n_rhs; }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::get_n_nonzero"
    int  Petsc_::get_n_nonzero() { return this->n_nz; }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::get_tolerence"
    double Petsc_::get_tolerence() { return (double)this->tol; }
  
  
#undef  __FUNCT__
#define __FUNCT__ "Petsc_::validate_sparse_matrix"
    bool Petsc_::validate_sparse_matrix() { return true; }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::init"
    int Petsc_::init(int argc, char **argv) {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;

      char help[] = "UHM interface to Petsc KSP solver\n";
      ierr = PetscInitialize(&argc, &argv, (char*)0, help);CHKERRQ(ierr);
      ierr = MPI_Comm_rank(PETSC_COMM_WORLD, &this->irank);CHKERRQ(ierr);
      ierr = MPI_Comm_size(PETSC_COMM_WORLD, &this->nrank);CHKERRQ(ierr);
      printf("[PETSC] Complete init ...\n");

#endif
      return true;
    }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::finalize"
    int Petsc_::finalize() {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;
      ierr = PetscFinalize();
      printf("[PETSC] Complete finalize ...\n");

#endif
      return true;
    }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::create_ksp"
    int Petsc_::create_ksp() {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;
    
      ierr = KSPCreate(PETSC_COMM_WORLD, &this->ksp);CHKERRQ(ierr);
      // - Options
      // SAME_PRECONDITIONER - Pmat is identical during successive linear solves
      // SAME_NONZERO_PATTERN - Pmat has same structure
      // DIFFERENT_NONZERO_PATTERN - different structure
      // 
      // Use same preconditioner for n_rhs
      ierr = KSPSetOperators(this->ksp, this->A, this->A, 
                             SAME_PRECONDITIONER);CHKERRQ(ierr);

      if (this->spd) {
        ierr = KSPSetType(this->ksp, KSPCG);CHKERRQ(ierr);
      } else {
        ierr = KSPSetType(this->ksp, KSPGMRES);CHKERRQ(ierr);
      }
      ierr = KSPSetTolerances(this->ksp,
                              this->get_tolerence(),
                              PETSC_DEFAULT,
                              PETSC_DEFAULT,
                              PETSC_DEFAULT);CHKERRQ(ierr);
      ierr = KSPSetInitialGuessNonzero(this->ksp, PETSC_TRUE);CHKERRQ(ierr);

#endif
    
      return true;
    }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::delete_ksp"
    int Petsc_::delete_ksp() {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;
    
      ierr = KSPDestroy(&this->ksp);CHKERRQ(ierr);

#endif

      return true;
    }

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::set_pc_on_ksp"
    int Petsc_::set_pc_on_ksp(char *type) {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;
      PC pc;

      // type 
      // PCNONE, PCJACOBI, PCSOR, PCLU, PCSHELL, PCBJACOBI, PCMG, PCILU
      // PCASM, PCKSP, PCCOMPOSITE, PCSPI, PCCHOLESKY, PCSAMG, PCHYPRE and so on
      // 
      // see petscpc.h
    
      ierr = KSPGetPC(this->ksp, &pc);CHKERRQ(ierr);
      ierr = PCSetType(pc, type);CHKERRQ(ierr);

#endif

      return true;
    }
  

#undef  __FUNCT__
#define __FUNCT__ "Petsc_::solve"
    int Petsc_::solve() {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;

      // n_rhs iteration
      for (int n=0;n<this->n_rhs;++n) {

        for (int k=0;k<this->n_dof;++k)
          ierr = VecSetValues(this->b, 
                              1, &k,
                              &(this->bb[n*this->n_dof+k]), 
                              INSERT_VALUES);CHKERRQ(ierr);

        ierr = VecAssemblyBegin(this->b);CHKERRQ(ierr);
        ierr = VecAssemblyEnd  (this->b);CHKERRQ(ierr);
      
        printf("[PETSC] DISPATCH VECTOR %d th b\n", n );

        ierr = KSPSetFromOptions(this->ksp);CHKERRQ(ierr);
        ierr = KSPSetUp(this->ksp);CHKERRQ(ierr);
        ierr = KSPSolve(this->ksp, this->b, this->x);CHKERRQ(ierr);
      
        ierr = KSPView(this->ksp, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

        PetscInt its;
        KSPConvergedReason reason;
        ierr = KSPGetConvergedReason(this->ksp, &reason);CHKERRQ(ierr);
        if (reason < 0) {
          printf("diverged reason = %d\n\n", reason);
          printf("KSP_DIVERGED_NULL                = -2\n");
          printf("KSP_DIVERGED_ITS                 = -3\n");
          printf("KSP_DIVERGED_DTOL                = -4\n");
          printf("KSP_DIVERGED_BREAKDOWN           = -5\n");
          printf("KSP_DIVERGED_BREAKDOWN_BICG      = -6\n");
          printf("KSP_DIVERGED_NONSYMMETRIC        = -7\n");
          printf("KSP_DIVERGED_INDEFINITE_PC       = -8\n");
          printf("KSP_DIVERGED_NAN                 = -9\n");
          printf("KSP_DIVERGED_INDEFINITE_MAT      = -10\n");
        } else {
          ierr = KSPGetIterationNumber(this->ksp, &its);
          printf("converged reason = %d, iteration = %d\n\n", reason, its);
          printf("KSP_CONVERGED_RTOL_NORMAL        =  1\n");
          printf("KSP_CONVERGED_ATOL_NORMAL        =  9\n");
          printf("KSP_CONVERGED_RTOL               =  2\n");
          printf("KSP_CONVERGED_ATOL               =  3\n");
          printf("KSP_CONVERGED_ITS                =  4\n");
          printf("KSP_CONVERGED_CG_NEG_CURVE       =  5\n");
          printf("KSP_CONVERGED_CG_CONSTRAINED     =  6\n");
          printf("KSP_CONVERGED_STEP_LENGTH        =  7\n");
          printf("KSP_CONVERGED_HAPPY_BREAKDOWN    =  8\n");
        }
      
        for (int k=0;k<this->n_dof;++k)
          ierr = VecGetValues(this->x,
                              1, &k,
                              &(this->xx[n*this->n_dof+k]));CHKERRQ(ierr);

      }

      for (int i = 0; i < this->show_n_rhs; ++i)
        printf("x [ %d ] = % E, b [ %d ] = % E\n",
               i, this->xx.at(i),
               i, this->bb.at(i) );

#endif

      return true;
    }


#undef  __FUNCT__
#define __FUNCT__ "Petsc_::check"
    int Petsc_::check() {

#ifdef UHM_INTERF_PETSC_ENABLE

      PetscErrorCode ierr;
      PetscReal norm;

      ierr = MatMult(this->A, this->x, this->r);
      ierr = VecAXPY(this->r, -1.0, this->b);
      ierr = VecNorm(this->r, NORM_2, &norm);CHKERRQ(ierr);
      ierr = PetscPrintf(PETSC_COMM_WORLD,"Norm of residual % E\n", norm);CHKERRQ(ierr);

#endif
      return true;
    }
  }
}
